'''
 * @Author: Benjay·Shaw
 * @Date: 2024-10-25 23:21:34
 * @LastEditors: Benjay·Shaw
 * @LastEditTime: 2024-10-29 19:44:23
 * @Description: 
'''
from utils import paddle_aux
import paddle


class eca_layer(paddle.nn.Layer):
    """Constructs a ECA module.
    Args:
        channel: Number of channels of the input feature map
        k_size: Adaptive selection of kernel size
    """

    def __init__(self, channel, k_size=3):
        super(eca_layer, self).__init__()
        self.avg_pool = paddle.nn.AdaptiveAvgPool2D(output_size=1)
        self.conv = paddle.nn.Conv1D(in_channels=1, out_channels=1,
            kernel_size=k_size, padding=(k_size - 1) // 2, bias_attr=False)
        self.sigmoid = paddle.nn.Sigmoid()

    def forward(self, x):
        y = self.avg_pool(x)
        y = self.conv(y.squeeze(axis=-1).transpose(perm=paddle_aux.
            transpose_aux_func(y.squeeze(axis=-1).ndim, -1, -2))).transpose(
            perm=paddle_aux.transpose_aux_func(self.conv(y.squeeze(axis=-1)
            .transpose(perm=paddle_aux.transpose_aux_func(y.squeeze(axis=-1
            ).ndim, -1, -2))).ndim, -1, -2)).unsqueeze(axis=-1)
        y = self.sigmoid(y)
        x = x * y.expand_as(y=x)
        return x * y.expand_as(y=x)
